// prefetcher.cc
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <csignal>
#include <ctime>
#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <cerrno>
#include <cstring>
#include <unordered_map>
#include <iomanip> // V2.1: 包含此头文件以使用 std::hex

extern "C" {
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "prefetcher.skel.h"
#include "prefetcher.h"
}

// ============================================================================
// V2: 用户态模式识别核心
// ============================================================================

#define PAGE_SIZE 4096
#define CONFIDENCE_THRESHOLD 3    // 连续多少次步长相同则认为模式稳定
#define PREFETCH_PAGES 16         // 每次预取多少页
#define MAX_REASONABLE_STRIDE (1024 * 1024 * 16) // 16MB，用于过滤异常步长

// 追踪每个线程（TID）的访问流状态
struct StreamTracker {
    uint32_t pid;
    uint32_t tid;
    unsigned long last_addr;
    long stride;
    int confidence;

    StreamTracker(uint32_t p, uint32_t t, unsigned long addr)
        : pid(p), tid(t), last_addr(addr), stride(0), confidence(1) {}
};

// 全局状态：使用TID作为key，追踪所有检测到的访问流
static std::unordered_map<uint32_t, StreamTracker> stream_trackers;

// ============================================================================
// 全局变量和辅助函数
// ============================================================================

std::ofstream log_file;
static volatile bool exiting = false;

// 统计计数器
static long faults_processed_count = 0;
static long pattern_detected_count = 0;
static long readahead_success_count = 0;
static long readahead_failure_count = 0;
static long map_failure_count = 0;

std::string get_timestamp() {
    time_t now = time(0);
    struct tm tstruct;
    char buf[80];
    tstruct = *localtime(&now);
    strftime(buf, sizeof(buf), "%Y-%m-%d %X", &tstruct);
    return buf;
}

void log_message(const std::string& msg) {
    std::string full_msg = "[" + get_timestamp() + "] " + msg;
    std::cout << full_msg << std::endl;
    if (log_file.is_open()) {
        log_file << full_msg << std::endl;
    }
}

// 地址转换函数 (保持不变)
bool find_file_offset(uint32_t pid, unsigned long vma_start, int& fd, off_t& file_offset, std::string& file_path) {
    std::string maps_path = "/proc/" + std::to_string(pid) + "/maps";
    std::ifstream maps_file(maps_path);
    if (!maps_file.is_open()) {
        return false;
    }
    std::string line;
    while (std::getline(maps_file, line)) {
        std::stringstream ss(line);
        std::string address_range, perms, offset_str, dev, inode, pathname;
        ss >> address_range >> perms >> offset_str >> dev >> inode;
        std::getline(ss, pathname);

        if (!pathname.empty() && pathname[0] == ' ') {
            pathname.erase(0, pathname.find_first_not_of(" \t"));
        }
        if (pathname.empty() || pathname[0] == '[') {
            continue;
        }

        unsigned long start, end, file_map_offset;
        sscanf(address_range.c_str(), "%lx-%lx", &start, &end);
        sscanf(offset_str.c_str(), "%lx", &file_map_offset);

        if (vma_start >= start && vma_start < end) {
            fd = open(pathname.c_str(), O_RDONLY);
            if (fd < 0) {
                log_message("  Error: Could not open file " + pathname + " (errno: " + std::to_string(errno) + ")");
                return false;
            }
            file_offset = (vma_start - start) + file_map_offset;
            file_path = pathname;
            return true;
        }
    }
    return false;
}

// Ring buffer 事件处理函数 (V2.1: 修正日志打印)
int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct fault_info *e = (struct fault_info*)data;
    faults_processed_count++;

    unsigned long current_addr = e->addr & ~(PAGE_SIZE - 1);

    auto it = stream_trackers.find(e->tid);

    // Case 1: 新的线程/访问流
    if (it == stream_trackers.end()) {
        stream_trackers.emplace(e->tid, StreamTracker(e->pid, e->tid, current_addr));
        // --- 修正点 ---
        std::stringstream log_msg;
        log_msg << "New stream detected for TID: " << e->tid << " at Addr: 0x" << std::hex << current_addr;
        log_message(log_msg.str());
        return 0;
    }

    // Case 2: 已存在的线程/访问流
    StreamTracker& tracker = it->second;
    long current_stride = current_addr - tracker.last_addr;
    tracker.last_addr = current_addr;

    // 过滤掉无效或异常的步长
    if (current_stride <= 0 || current_stride > MAX_REASONABLE_STRIDE) {
        tracker.confidence = 1;
        tracker.stride = 0; // 重置步长
        log_message("TID " + std::to_string(e->tid) + ": Invalid stride detected (" + std::to_string(current_stride) + "). Resetting confidence.");
        return 0;
    }

    // 更新置信度
    if (current_stride == tracker.stride) {
        tracker.confidence++;
    } else {
        tracker.stride = current_stride;
        tracker.confidence = 1;
    }

    // --- 修正点 ---
    std::stringstream log_msg;
    log_msg << "TID " << e->tid << ": Fault at 0x" << std::hex << current_addr
            << ", Stride: " << std::dec << current_stride
            << ", Confidence: " << tracker.confidence;
    log_message(log_msg.str());

    // 判断是否达到预取阈值
    if (tracker.confidence >= CONFIDENCE_THRESHOLD) {
        pattern_detected_count++;
        tracker.confidence = 0; // 成功预测后重置，防止对同一流的每次缺页都预取

        unsigned long prefetch_addr = current_addr + tracker.stride;
        // 自适应预取大小：至少预取16页，但如果步长很大，就预取2倍步长
        size_t prefetch_size = std::max((size_t)PREFETCH_PAGES * PAGE_SIZE, (size_t)2 * (size_t)tracker.stride);

        // --- 修正点 ---
        log_msg.str(""); // 清空 stringstream
        log_msg.clear(); // 清除状态标志
        log_msg << "  -> STEADY PATTERN DETECTED for TID " << e->tid
                << ". Prefetching " << prefetch_size / 1024 << " KB"
                << " from predicted addr 0x" << std::hex << prefetch_addr;
        log_message(log_msg.str());

        int fd = -1;
        off_t file_offset = 0;
        std::string file_path;
        if (find_file_offset(e->pid, prefetch_addr, fd, file_offset, file_path)) {
            log_message("  -> Mapped to file: " + file_path + ", FD: " + std::to_string(fd) + ", Offset: " + std::to_string(file_offset));

            long ret = syscall(SYS_readahead, fd, file_offset, prefetch_size);

            if (ret == 0) {
                log_message("  -> readahead() syscall SUCCEEDED.");
                readahead_success_count++;
            } else {
                log_message("  -> readahead() syscall FAILED. Return code: " + std::to_string(ret) + ", errno: " + std::to_string(errno) + " (" + strerror(errno) + ")");
                readahead_failure_count++;
            }
            close(fd);
        } else {
            log_message("  -> Failed to map address to file offset. Prediction may be out of bounds.");
            map_failure_count++;
        }
    }

    return 0;
}


int print_libbpf_log(enum libbpf_print_level level, const char *format, va_list args) {
    if (level == LIBBPF_WARN) { // 只打印警告和错误
        char buffer[1024];
        vsnprintf(buffer, sizeof(buffer), format, args);
        log_message("libbpf: " + std::string(buffer));
    }
    return 0;
}

static void sig_handler(int sig) {
    exiting = true;
}

// ============================================================================
// Main 函数
// ============================================================================
int main(int argc, char **argv) {
    if (argc > 1 && (std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help")) {
        std::cout << "Usage: " << argv[0] << " [target_pid]" << std::endl;
        return 0;
    }

    log_file.open("prefetcher.log", std::ios::out | std::ios::trunc);
    log_message("Prefetcher V2 starting...");

    libbpf_set_print(print_libbpf_log);

    struct prefetcher_bpf *skel = prefetcher_bpf__open();
    if (!skel) {
        log_message("Error: Failed to open BPF skeleton");
        return 1;
    }

    int err = prefetcher_bpf__load(skel);
    if (err) {
        log_message("Error: Failed to load BPF skeleton. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }

    // 设置目标PID
    struct config cfg = {};
    if (argc > 1) {
        cfg.target_pid = std::stoi(argv[1]);
        log_message("Targeting PID: " + std::string(argv[1]));
    } else {
        cfg.target_pid = 0;
        log_message("Monitoring all processes.");
    }
    int map_fd = bpf_map__fd(skel->maps.config_map);
    const int key = 0;
    err = bpf_map_update_elem(map_fd, &key, &cfg, BPF_ANY);
    if (err) {
        log_message("Error: Failed to update config map. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }

    err = prefetcher_bpf__attach(skel);
    if (err) {
        log_message("Error: Failed to attach BPF skeleton. Errno: " + std::to_string(err) + ", Message: " + strerror(-err));
        prefetcher_bpf__destroy(skel);
        return 1;
    }
    log_message("eBPF program loaded and attached successfully.");

    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    struct ring_buffer *rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event, NULL, NULL);
    if (!rb) {
        log_message("Error: Failed to create ring buffer");
        err = -1;
        goto cleanup;
    }

    log_message("Waiting for events... Press Ctrl+C to exit.");
    while (!exiting) {
        err = ring_buffer__poll(rb, 100 /* timeout, ms */);
        if (err == -EINTR) {
            err = 0;
            break;
        }
        if (err < 0) {
            log_message("Error polling ring buffer: " + std::to_string(err));
            break;
        }
    }

cleanup:
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);

    log_message("Prefetcher stopped. Final statistics:");
    std::stringstream final_stats;
    final_stats << "  - Total page faults processed: " << faults_processed_count << std::endl;
    final_stats << "  - Steady patterns detected: " << pattern_detected_count << std::endl;
    final_stats << "  - Successful readahead calls: " << readahead_success_count << std::endl;
    final_stats << "  - Failed readahead calls: " << readahead_failure_count << std::endl;
    final_stats << "  - Address mapping failures: " << map_failure_count;
    log_message(final_stats.str());

    if (log_file.is_open()) {
        log_file.close();
    }
    return err < 0 ? -err : 0;
}